#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Virtual machine image builder
"""

import os
import subprocess
import tempfile
import shutil
import logging
import platform
import requests
from tqdm import tqdm
from typing import Dict, Any, Optional
from trustimage.core.parser import CommandType, TrustedImageParser
from trustimage.core.dm_verity import DMVerity
from trustimage.core.kernel_cmdline import KernelCmdline

logger = logging.getLogger(__name__)

class BuildError(Exception):
    """Error during build process"""
    pass

class ImageBuilder:
    """Virtual machine image builder"""
    
    def __init__(self, parser: TrustedImageParser, output_dir: str = '.'):
        """
        Initialize builder
        
        Args:
            parser: TrustedImageParser instance
            output_dir: Output directory
        """
        self.parser = parser
        self.output_dir = output_dir
        self.temp_dir: Optional[str] = None
        self.commands = []
        self.base_image: Optional[str] = None
        self.output_format: Optional[str] = None
        self.output_name: Optional[str] = None
        self.current_image: Optional[str] = None
        
        # Check if virt-customize is installed
        self._check_dependencies()
    
    def _check_dependencies(self):
        """Check if required tools are installed"""
        try:
            subprocess.run(['virt-customize', '--version'], 
                         stdout=subprocess.PIPE, 
                         stderr=subprocess.PIPE, 
                         check=True)
            logger.debug("virt-customize found")
        except (subprocess.SubprocessError, FileNotFoundError):
            raise BuildError("virt-customize tool not found, please install libguestfs-tools package")
    
    def _get_system_architecture(self) -> str:
        """获取系统架构"""
        arch_map = {
            'x86_64': 'amd64',
            'aarch64': 'arm64',
            'armv7l': 'armhf',
        }
        return arch_map.get(platform.machine(), 'amd64')

    def _parse_image_name(self, image_name: str) -> tuple:
        """
        解析镜像名称
        
        Args:
            image_name: 镜像名称，格式如 ubuntu:24.10 或 /path/to/local/image.img
        
        Returns:
            tuple: (distro, version) 或 (None, None) 如果是本地路径
        """
        if os.path.exists(image_name) or image_name.startswith('/'):
            return None, None
        
        if ':' in image_name:
            distro, version = image_name.split(':', 1)
            return distro.lower(), version
        
        return None, None

    def _build_download_url(self, distro: str, version: str) -> str:
        """
        根据发行版类型构建下载URL (switch-like 逻辑)
        
        Args:
            distro: 发行版名称
            version: 版本号
        
        Returns:
            str: 下载URL
        """
        arch = self._get_system_architecture()
        
        if distro == 'ubuntu':
            return f"https://cloud-images.ubuntu.com/releases/{version}/release/ubuntu-{version}-server-cloudimg-{arch}.img"
        elif distro == 'centos':
            # CentOS support not implemented yet
            raise BuildError("CentOS distribution is not supported yet")
        elif distro == 'openEuler':
            # openEuler support not implemented yet
            raise BuildError("openEuler distribution is not supported yet")
        else:
            raise BuildError(f"Unsupported distribution: {distro}")

    def _get_cache_path(self, distro: str, version: str) -> str:
        """
        获取缓存文件路径
        
        Args:
            distro: 发行版名称
            version: 版本号
        
        Returns:
            str: 本地缓存文件路径
        """
        cache_dir = os.path.expanduser('~/.trustimage/cache')
        os.makedirs(cache_dir, exist_ok=True)
        
        arch = self._get_system_architecture()
        
        if distro == 'ubuntu':
            filename = f"ubuntu-{version}-server-cloudimg-{arch}.img"
        elif distro == 'centos':
            filename = f"CentOS-{version}-{arch}-GenericCloud.qcow2"
        elif distro == 'openEuler':
            filename = f"openEuler-{version}-LTS-{arch}.qcow2.xz"
        else:
            filename = f"{distro}-{version}-{arch}.img"
        
        return os.path.join(cache_dir, distro, version, filename)

    def _download_image(self, url: str, local_path: str) -> bool:
        """
        下载镜像文件
        
        Args:
            url: 下载URL
            local_path: 本地保存路径
        
        Returns:
            bool: 下载是否成功
        """
        try:
            # 创建目标目录
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            
            logger.info(f"Downloading image from: {url}")
            
            response = requests.get(url, stream=True)
            response.raise_for_status()
            
            total_size = int(response.headers.get('content-length', 0))
            
            with open(local_path, 'wb') as f, tqdm(
                desc="Downloading",
                total=total_size,
                unit='B',
                unit_scale=True,
                unit_divisor=1024,
            ) as pbar:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        pbar.update(len(chunk))
            
            logger.info(f"Image downloaded successfully: {local_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to download image: {e}")
            if os.path.exists(local_path):
                os.remove(local_path)
            return False

    def _resolve_base_image(self, image_name: str) -> str:
        """
        解析并获取基础镜像，支持本地文件和远程下载
        
        Args:
            image_name: 镜像名称或路径
        
        Returns:
            str: 本地镜像文件路径
        """
        # 检查是否为本地文件路径
        if os.path.exists(image_name):
            logger.info(f"Using local image: {image_name}")
            return image_name
        
        # 解析镜像名称
        distro, version = self._parse_image_name(image_name)
        
        if distro is None:
            # 如果无法解析且不是本地文件，抛出错误
            raise BuildError(f"Invalid image name or file not found: {image_name}")
        
        # 检查本地缓存
        cache_path = self._get_cache_path(distro, version)
        if os.path.exists(cache_path):
            logger.info(f"Using cached image: {cache_path}")
            return cache_path
        
        # 从远程下载
        try:
            download_url = self._build_download_url(distro, version)
            if self._download_image(download_url, cache_path):
                return cache_path
            else:
                raise BuildError(f"Failed to download image: {image_name}")
        except Exception as e:
            raise BuildError(f"Failed to resolve image {image_name}: {e}")
    
    def prepare(self):
        """Prepare build environment"""
        # Parse configuration file
        self.commands = self.parser.parse()
        base_image_name = self.parser.get_base_image()
        self.output_format, self.output_name = self.parser.get_output_info()
        
        if not base_image_name or not self.output_name:
            raise BuildError("Base image or output name not specified")
        
        # 解析并获取基础镜像（支持本地和远程）
        self.base_image = self._resolve_base_image(base_image_name)
        
        # Create temporary directory
        self.temp_dir = tempfile.mkdtemp(prefix='trustimage-')
        logger.debug(f"Created temporary directory: {self.temp_dir}")
        
        # Copy base image to temporary directory
        self.current_image = os.path.join(self.temp_dir, os.path.basename(self.base_image))
        try:
            shutil.copy2(self.base_image, self.current_image)
            logger.info(f"Copied base image to: {self.current_image}")
        except FileNotFoundError:
            raise BuildError(f"Base image does not exist: {self.base_image}")
    
    def build(self) -> str:
        """
        Build image
        
        Returns:
            str: Output image path
        """
        if not self.commands:
            raise BuildError("Configuration file not parsed, please call prepare() first")
        
        if not self.current_image or not self.output_name:
            raise BuildError("Current image or output name not set")
        
        logger.info("Starting image build")
        
        # Process each command in order from configuration file
        for cmd in self.commands:
            if not self._execute_command(cmd):
                raise BuildError(f"Command execution failed: {cmd}")
        
        # Move final image to output directory
        output_path = os.path.join(self.output_dir, f"{self.output_name}.{self.output_format}")
        shutil.move(self.current_image, output_path)
        
        logger.info(f"Image build completed: {output_path}")
        return output_path
    
    def _execute_command(self, cmd) -> bool:
        """
        Execute single command
        
        Args:
            cmd: Command object
        
        Returns:
            bool: Execution success or not
        """
        logger.debug(f"Executing command: {cmd}")
        
        try:
            
            if cmd.type == CommandType.FROM:
                # FROM command handled in prepare stage
                return True
            elif cmd.type == CommandType.INSTALL:
                return self._handle_install(cmd.args)
            elif cmd.type == CommandType.RUN:
                return self._handle_run(cmd.args)
            elif cmd.type == CommandType.COPY:
                return self._handle_copy(cmd.args)
            elif cmd.type == CommandType.DM_VERITY:
                return self._handle_dm_verity(cmd.args)
            elif cmd.type == CommandType.SET_KERNEL_CMDLINE:
                return self._handle_kernel_cmdline(cmd.args)
            elif cmd.type == CommandType.SET_ROOTPASSWD:
                return self._handle_root_password(cmd.args)
            elif cmd.type == CommandType.OUTPUT:
                # OUTPUT command handled in prepare stage
                return True
            else:
                logger.warning(f"Unhandled command type: {cmd.type}")
                return False
                
        except Exception as e:
            logger.error(f"Command execution failed: {e}")
            return False
    
    def _handle_install(self, args) -> bool:
        """Handle INSTALL command"""
        if not args:
            logger.warning("INSTALL command missing arguments")
            return False
        
        if not self.current_image:
            raise BuildError("Current image not set")
        
        try:
            cmd = ['virt-customize', '-a', self.current_image, '--install', args[0]]
            subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to install package: {e.stderr}")
            return False
    
    def _handle_run(self, args) -> bool:
        """Handle RUN command"""
        if not args:
            logger.warning("RUN command missing arguments")
            return False
        
        if not self.current_image:
            raise BuildError("Current image not set")
        
        try:
            cmd = ['virt-customize', '-a', self.current_image, '--run-command', ' '.join(args)]
            subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to execute command: {e.stderr}")
            return False
    
    def _handle_copy(self, args) -> bool:
        """Handle COPY command"""
        if len(args) < 2:
            logger.warning("COPY command requires source file and destination path")
            return False
        
        if not self.current_image:
            raise BuildError("Current image not set")
        
        src, dst = args[0], args[1]
        if not os.path.exists(src):
            logger.warning(f"Source file does not exist: {src}")
            return False
            
        # Ensure destination path is absolute
        if not dst.startswith('/'):
            dst = '/' + dst
        
        try:
            # Handle source file path
            if not os.path.isabs(src):
                src = os.path.abspath(src)
            
            # Determine source file type and choose command
            if os.path.isdir(src):
                # If directory, use --copy-in command
                cmd = ['virt-customize', '-a', self.current_image, '--copy-in', f'{src}:{dst}']
            else:
                # If file, use --upload command
                cmd = ['virt-customize', '-a', self.current_image, '--upload', f'{src}:{dst}']
            
            logger.debug(f"Execute copy command: {' '.join(cmd)}")
            subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to copy file: {e.stderr}")
            return False
    
    def _handle_dm_verity(self, args) -> bool:
        """Handle DM_VERITY command"""
        if not args:
            logger.warning("DM_VERITY command missing arguments")
            return False
        
        if not self.current_image:
            raise BuildError("Current image not set")
        
        try:
            # Create DMVerity instance
            dm_verity = DMVerity()
            
            # Parse parameters
            fs_type = args[0] if args else 'rootfs'
            salt = args[1].split('=')[1] if len(args) > 1 and '=' in args[1] else None
            
            # Set DM_VERITY
            verity_info = dm_verity.setup(self.current_image, fs_type, salt)
            
            
            return True
            
        except Exception as e:
            logger.error(f"DM_VERITY setup failed: {e}")
            return False
    
    def _handle_kernel_cmdline(self, args) -> bool:
        """
        Handle SET_KERNEL_CMDLINE command
        
        Args:
            args: List of command arguments
        
        Returns:
            bool: Success or not
        """
        if not args:
            logger.warning("SET_KERNEL_CMDLINE command missing arguments")
            return False
        
        if not self.current_image:
            raise BuildError("Current image not set")
        
        try:
            # Create KernelCmdline instance
            kernel_cmdline = KernelCmdline(self.current_image)
            
            # Determine replace or update mode
            if args[0].lower() == 'replace':
                # Replace mode: use remaining args as new full command line
                new_cmdline = ' '.join(args[1:])
                return kernel_cmdline.replace_cmdline(new_cmdline)
            else:
                # Update mode: use all args to update existing command line
                return kernel_cmdline.update_cmdline(args)
            
        except Exception as e:
            logger.error(f"Failed to set kernel command line: {e}")
            return False
    
    def _handle_root_password(self, args) -> bool:
        """
        Handle SET_ROOTPASSWD command
        
        Args:
            args: List of command arguments, should contain root password
        
        Returns:
            bool: Success or not
        """
        if not args:
            logger.warning("SET_ROOTPASSWD command missing password argument")
            return False
        
        if not self.current_image:
            raise BuildError("Current image not set")
        
        try:
            password = args[0]
            cmd = ['virt-customize', '-a', self.current_image, '--root-password', f'password:{password}']
            subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            logger.info("Root password set successfully")
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to set root password: {e.stderr}")
            return False
    
    def cleanup(self):
        """Clean up temporary files"""
        if self.temp_dir and os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)
            self.temp_dir = None 